package matlab;

import java.io.ByteArrayOutputStream;
import java.io.DataOutputStream;
import java.io.File;
import java.io.FileNotFoundException;
import java.io.FileOutputStream;
import java.io.IOException;
import java.util.HashMap;
import java.util.Map;
import java.util.Set;
import java.util.Map.Entry;

import java.io.FileInputStream;

/**
 * Class for handling MAT-file version 5.0 (matlab 6.x ->)
 * @author risto
 *
 */
public class Matlab5File {
	
	// File header
	Header hdr;
	Map<String, IVariable> workspace;

	/**
	 * Get variable.
	 * @param varName
	 * @return
	 */
	public IVariable get(String varName) { return workspace.get(varName); }
	
	/**
	 * List of variable names
	 * @return
	 */
	public Set<String> listVariables()
	{
		return workspace.keySet();
	}

	public void addDouble(String name, double value)
	{
		workspace.put(name, new DoubleVariable(name, value));
	}

	public void addDoubleArray(String name, double[] value, int m, int n)
	{
		workspace.put(name, new DoubleArrayVariable(name, value, m, n));
	}

	public void addText(String name, String value)
	{
		workspace.put(name, new TextVariable(name, value));
	}

	/**
	 * Load existing file
	 * @param f
	 * @throws FileNotFoundException
	 * @throws IOException
	 */
	public Matlab5File(File f) throws FileNotFoundException, IOException {

		byte[] data = fileBytes(f);
		
		// read header
		hdr = new Header(data);
		System.out.println(hdr.description);

		// read contents
		int bytesLeft = data.length - 128;
		int offset = 128;
		workspace = new HashMap<String, IVariable>();
		while(bytesLeft > 0) {
			Variable var = new Variable(data, offset);
			workspace.put(var.getName(), var);
			offset += var.paddedLength();
			bytesLeft -= var.paddedLength();
		}
	}

	/** 
	 * Create new file 
	 */
	public Matlab5File() {
		hdr = new Header();
		workspace = new HashMap<String, IVariable>();
	}

	/**
	 * Write to file.
	 * @param f
	 * @throws IOException
	 */
	public void write(File f) throws IOException {
		FileOutputStream fos = new FileOutputStream(f);
		DataOutputStream dos = new DataOutputStream(fos);
		hdr.writeTo(dos); // Write header
		for (Entry<String,IVariable> e : workspace.entrySet()) {
			writeVariable(dos, e.getValue()); // Write all variables
		}
		dos.flush();
		fos.close();
	}
	
	private int padLen(int l)
	{
		if (l%8 == 0) return 0;
		return 8 - (l%8);
	}
	
	private void pad(DataOutputStream os, int pad) throws IOException
	{
		for (int i = 0; i < pad; i++) os.write((byte)0x00);
	}
	
	
	private void writeVariable(DataOutputStream os, IVariable var) throws IOException
	{
		ByteArrayOutputStream bos;
		DataOutputStream dos;
		
		byte arrayFlags[], dimensions[], arrayName[], real[]; 
		
		// Array Flags
		bos = new ByteArrayOutputStream();
		dos = new DataOutputStream(bos);
		dos.writeInt(ChunkType.UINT32.value);  // Uint32
		dos.writeInt(8); // 8 bytes
		dos.writeInt(var.getType().value); // array type
		dos.writeInt(0x00); // undefined
		dos.flush();
		arrayFlags = bos.toByteArray();
		
		// Dimensions
		bos = new ByteArrayOutputStream();
		dos = new DataOutputStream(bos);
		int[] dim = var.getDimensions();
		int dataN = 1;
		if (dim == null) dim = new int[] { 1, 1 };
		for (int d : dim) dataN *= d;
		dos.writeInt(ChunkType.INT32.value);  // int32
		dos.writeInt(dim.length * 4); // 4bytes per dimension
		for (int i = 0; i < dim.length; i++) dos.writeInt(dim[i]);
		dos.flush();
		dimensions = bos.toByteArray();

		// ArrayName
		bos = new ByteArrayOutputStream();
		dos = new DataOutputStream(bos);
		String name = var.getName();
		if (name.length() <= 4) { // Short ?
			dos.writeShort(name.length());
			dos.writeShort(ChunkType.INT8.value);
			dos.write(name.getBytes());
			pad(dos, 4 - name.getBytes().length);
		}
		else { // Long
			dos.writeInt(ChunkType.INT8.value);
			dos.writeInt(name.getBytes().length);
			dos.write(name.getBytes());
			pad(dos, padLen(name.getBytes().length));
		}
		dos.flush();
		arrayName = bos.toByteArray();

		// Real part
		bos = new ByteArrayOutputStream();
		dos = new DataOutputStream(bos);
		
		// Double, always long, no need to pad
		if (var.isDouble()) {
			double[] arr = var.getDoubleArray();
			dos.writeInt(ChunkType.DOUBLE.value);
			dos.writeInt(8*arr.length);
			for (double d : arr) dos.writeDouble(d);
		}
		
		else if (var.isText()) { // write as uint16
			dos.writeInt(ChunkType.UINT16.value);
			String t = var.getText();
			dos.writeInt(2*t.length());
			for (char c : t.toCharArray()) {
				dos.writeShort(c);
			}
			pad(dos, padLen(2*t.length()));
		}
		
		else if (var.isInteger()) { // write as UINT32
			int[] arr = var.getIntegerArray();
			dos.writeInt(ChunkType.UINT32.value);
			dos.writeInt(4*arr.length);
			for (int i : arr) dos.writeInt(i);
			pad(dos, padLen(4*arr.length));
		}
		dos.flush();
		real = bos.toByteArray();

		// Matrix data
		os.writeInt(ChunkType.MATRIX.value);
		int totalSize = arrayFlags.length + dimensions.length + arrayName.length + real.length;
		os.writeInt(totalSize);
		os.write(arrayFlags);
		os.write(dimensions);
		os.write(arrayName);
		os.write(real);
		System.out.println("WRITE " + var.getName());
		System.out.println("SIZE " + totalSize);
		System.out.println("flags " + arrayFlags.length + " bytes");
		System.out.println("dim " + dimensions.length + " bytes");
		System.out.println("name " + arrayName.length + " bytes");
		System.out.println("real " + real.length + " bytes");
		os.flush();
	}
	
	/**
	 * @param args
	 */
	public static void main(String[] args) throws Exception
	{
		File f = new File("test.mat");
		Matlab5File m5 = new Matlab5File(f);
		System.out.println(m5.listVariables());
		for (String var : m5.listVariables()) {
			System.out.println(m5.get(var));
		}
		f = new File("test_out.mat");
		m5.write(f);
		
		f = new File("test_out2.mat");
		Matlab5File newFile = new Matlab5File();
		newFile.addText("hello", "Hello, world!");
		newFile.addDouble("n", 123);
		newFile.addDouble("pi_est", 3.14159264);
		newFile.addDoubleArray("matrix_ident", new double[] { 1, 0, 0, 0, 1, 0, 0, 0, 1 }, 3, 3);
		newFile.addDoubleArray("some_points", DoubleArrayVariable.transpose(new double[] { 
				1.0, 1.0, 1.0, 
				1.0, 2.0, 1.0, 
				1.0, 3.0, 1.0, 
				1.0, 4.0, 1.0,
				2.0, 1.0, 1.0,
				3.0, 1.0, 1.0,
				2.0, 2.0, 1.0 
				}, 3, 7),
				7, 3);
		newFile.write(f);
	}

	
    /**
     *
     * Read contents of File f into byte array.
     *
     * @param f - file to read
     * @return byte[] array of file contents.
     * @throws FileNotFoundException
     * @throws IOException
     */
    public static byte[] fileBytes(File f) throws FileNotFoundException, IOException
    {
        byte data[] = new byte[(int)f.length()];
        FileInputStream fis = new FileInputStream(f);
        int toRead = (int) f.length();
        int r = 0;
        while(toRead > 0) {
            r = fis.read(data);
            if (r == -1) break;
            toRead -= r;
        }
        if (toRead != 0) throw new IOException("Partial read of " + f);
        fis.close();
        return data;
    }
}
